"""
模型的评估
"""
import torch
import numpy as np
from tqdm import tqdm

import config
from dnn.sort.siamese import SiameseNetWork
from dnn.sort.dataset import data_loader as test_data


# 模型的创建和加载
model = SiameseNetWork().to(config.device)
model.load_state_dict(torch.load(config.sort_save_model_path))


def evaluate():
    model.eval()
    bar = tqdm(enumerate(test_data), total=len(test_data), ascii=True, desc="测试")
    loss_list = []
    acc_list = []
    for idx, (input1, input2, target) in bar:
        input1 = input1.to(config.device)
        input2 = input2.to(config.device)
        target = target.to(config.device)
        output = model(input1, input2)
        loss = torch.nn.functional.nll_loss(output, target)
        loss_list.append(loss.item())
        pre = torch.max(output, dim=-1)[-1]
        # print(target.size())
        # print(pre.size())
        cur_acc = pre.eq(target).float().mean()
        acc_list.append(cur_acc.cpu().item())

    print(np.mean(loss_list), np.mean(acc_list))

    pass

